#include "core/main.h"

/*
Two types of questions to ask. We chould use different variable sets for each question, although for simplicity we unify them because there is lots of overlap.
There is no harm in having extra always-constant variables lying around, they will never be used as a split variable.

A unique sprite is a sprite that has been observed at most once per frame.

*** SPRITE-AT QUESTION ***
Is there a sprite of type Z at coordinate X,Y

*** VARIABLE CHANGE QUESTION ***
Absolute: Should the value of variable Z be set to X?
Relative: Should the value of variable Z be incremented by X?

Variables:
2: X, Y (meaningless for variable change question)
Per-frame information (historyFrameCount):
  (stateWindowSize * 2 + 1)^2: window centered at X, Y (meaningless for variable change question)
  buttonCount: controller state
  N: all observable integers
  S: count of all sprites
  2 + (stateWindowSize * 2 + 1)^2: for each unique sprite, record absolute X, Y starting at first instance of each sprite (do not deal with overlapping sprites because it's not a common case)
*/


void ModelLearner::LearnModel(const ReplayCollection &replays, const GameModelParameters &parameters, GameModel &model)
{
    _parameters = parameters;
    _game = replays.Replays()[0]->SourceGame();

    model._game = _game;

    InitFrames(replays);
    InitSpriteInstances();
    MakeGeneralFeatures();
    TrainDecisionTrees(model);

    model._parameters = _parameters;
}

void ModelLearner::InitFrames(const ReplayCollection &replays)
{
    _parameters.totalSpriteCount = _game->Collection().NonFontSpriteCount();

    vector<int> spriteMaxCounts(_parameters.totalSpriteCount, 0);

    _parameters.dimensions = Vec2i(1, 1);

    _frames.resize(replays.TotalFrameCount());
    int learnerFrameIndex = 0;
    for(UINT replayIndex = 0; replayIndex < replays.Replays().size(); replayIndex++)
    {
        const Replay &curReplay = *replays.Replays()[replayIndex];
        for(UINT replayFrameIndex = 0; replayFrameIndex < curReplay.Frames().size(); replayFrameIndex++)
        {
            const Frame &curFrame = curReplay.Frames()[replayFrameIndex];
            _frames[learnerFrameIndex] = unique_ptr<GameModelFrame>(new GameModelFrame(&curFrame));
            GameModelFrame &learnerFrame = *_frames[learnerFrameIndex++];

            learnerFrame.ComputeSpriteCounts(_parameters);
            for(UINT spriteID = 0; spriteID < _parameters.totalSpriteCount; spriteID++)
            {
                spriteMaxCounts[spriteID] = Math::Max(spriteMaxCounts[spriteID], learnerFrame.spriteCount[spriteID]);
            }

            for(UINT spriteIndex = 0; spriteIndex < curFrame.sprites.size(); spriteIndex++)
            {
                const SpriteLocation &curSprite = curFrame.sprites[spriteIndex];
                if(curSprite.spriteID < _parameters.totalSpriteCount)
                {
                    _parameters.dimensions = Vec2i::Maximize(_parameters.dimensions, curFrame.sprites[spriteIndex].coord);
                }
            }
        }
    }

    //_frames.resize(1000);

    _parameters.uniqueSprites.clear();
    for(UINT spriteIndex = 0; spriteIndex < _parameters.totalSpriteCount; spriteIndex++)
    {
        if(spriteMaxCounts[spriteIndex] == 1)
        {
            _parameters.uniqueSprites.push_back(spriteIndex);
        }
    }
}

void ModelLearner::InitSpriteInstances()
{
    const UINT startSkipCount = 5;

    _spriteInstances.resize(_parameters.totalSpriteCount);
    for(UINT frameIndex = _parameters.historyFrameCount + startSkipCount; frameIndex < _frames.size(); frameIndex++)
    {
        GameModelFrame &curFrame = *_frames[frameIndex];
        for(UINT spriteIndex = 0; spriteIndex < curFrame.sourceFrame->sprites.size(); spriteIndex++)
        {
            int spriteID = curFrame.sourceFrame->sprites[spriteIndex].spriteID;
            if(spriteID < int(_spriteInstances.size()))
            {
                _spriteInstances[spriteID].push_back(LearnerSprite(frameIndex, spriteIndex));
            }
        }
    }
}

void ModelLearner::MakeGeneralFeatures()
{
    const UINT stateWindowVariableCount = 2 + Math::Square(_parameters.stateWindowSize * 2 + 1); // the "2 +" records the absolute x and y coordinates of the window
    _parameters.observedIntegerCount = (*_frames[0]).sourceFrame->observedInts.size();

    _parameters.perFrameVariableCount = 
        Constants::buttonCount +               // state of all buttons
        _parameters.observedIntegerCount +     // state of all observed integers
        _parameters.totalSpriteCount +         // global sprite count
        stateWindowVariableCount +             // this sprite window
        stateWindowVariableCount * _parameters.uniqueSprites.size(); // unique sprite windows

    _parameters.totalVariableCount = _parameters.perFrameVariableCount * _parameters.historyFrameCount;

    cout << "Making all feature vectors..." << endl;
    cout << _parameters.perFrameVariableCount << " variables, " << _frames.size() << " frames, " << (_parameters.perFrameVariableCount * sizeof(short) * _frames.size()) / 1024.0 / 1024.0 << "MB" << endl;

    for(int frameIndex = 0; frameIndex < int(_frames.size()); frameIndex++)
    {
        GameModelFrame &curFrame = *_frames[frameIndex];
        curFrame.historyFrames.resize(_parameters.historyFrameCount);
        for(int frameHistoryIndex = 0; frameHistoryIndex < _parameters.historyFrameCount; frameHistoryIndex++)
        {
            curFrame.historyFrames[frameHistoryIndex] = _frames[Math::Bound(frameIndex - frameHistoryIndex, 0, int(_frames.size()))].get();
        }
        curFrame.cachedFeatures = new short[_parameters.totalVariableCount];
        curFrame.MakeFeatureVector(_parameters, -9999, -9999, curFrame.cachedFeatures);
    }
}

void ModelLearner::InitializeDatasets()
{
    const UINT frameCount = _frames.size();
    
    _variableDataset.examples.resize(frameCount - 1);
    for(UINT frameIndex = 0; frameIndex < frameCount - 1; frameIndex++)
    {
        LearnerExample &curExample = _variableDataset.examples[frameIndex];
        GameModelFrame &curFrame = *_frames[frameIndex];
        GameModelFrame &nextFrame = *_frames[frameIndex + 1];

        curExample.input.data = curFrame.cachedFeatures;
        curExample.input.length = _parameters.totalVariableCount;
    }

    _spriteDataset.examples.resize(_parameters.spriteDatasetSize);
    for(int exampleIndex = 0; exampleIndex < _parameters.spriteDatasetSize; exampleIndex++)
    {
        LearnerExample &curExample = _spriteDataset.examples[exampleIndex];
        curExample.input.data = new short[_parameters.totalVariableCount];
        curExample.input.length = _parameters.totalVariableCount;
    }
}

void ModelLearner::TrainDecisionTrees(GameModel &model)
{
    InitializeDatasets();

    TrainSpriteTrees(model);
    TrainVariableTrees(model);
}

inline void writeErrorDetails(const string &message, const ClassificationErrorDetails &errorDetails, ostream &stream)
{
    stream << message << " error: " << errorDetails.errorRate << endl;
    stream << message << " precision: " << errorDetails.precision << endl;
    stream << message << " recall: " << errorDetails.precision << endl;
}

void ModelLearner::TrainVariableTrees(GameModel &model)
{
    const UINT frameCount = _frames.size();
    
    ofstream logFile("variableTrainingLog.txt");
    ComposeStream stream(logFile, cout);

    ofstream trainingFile("variableTrainingSummary.txt");
    trainingFile << "Variable name\tAvg. training error\tAvg. test error" << endl;

    DecisionTreeConfig config(50, 1, 100);

    model._variableModels.resize(_parameters.observedIntegerCount);

    //
    // Learn decision trees for all variables
    //
    for(UINT variableIndex = 0; variableIndex < _parameters.observedIntegerCount; variableIndex++)
    {
        double averageTestError = 0.0;
        double averageTrainingError = 0.0;

        //
        // Decide what kind of questions to ask
        //
        set<int> absolutes, relatives;
        ComputeVariableChanges(variableIndex, absolutes, relatives);

        stream << "Constructing trees for variable " << variableIndex << " " << _game->ObservedIntNames()[variableIndex] << endl;

        GameVariableModel *variableModel = new GameVariableModel;
        model._variableModels[variableIndex] = variableModel;
        variableModel->trees.resize(absolutes.size() + relatives.size());
        int treeIndex = 0;

        //
        // Compute absolute trees
        //
        for(auto absoluteIterator = absolutes.begin(); absoluteIterator != absolutes.end(); absoluteIterator++)
        {
            for(UINT frameIndex = 0; frameIndex < frameCount - 1; frameIndex++)
            {
                _variableDataset.examples[frameIndex].output = Utility::BoolToInt(_frames[frameIndex + 1]->sourceFrame->observedInts[variableIndex] == *absoluteIterator);
            }

            LearnerDataset trainingSet, testSet;
            _variableDataset.PartitionDataset(trainingSet, testSet, 0.7);

            stream << "Training tree for absolute value: " << *absoluteIterator << endl;

            GameVariableTree &variableTree = variableModel->trees[treeIndex++];
            variableTree.absolute = true;
            variableTree.value = *absoluteIterator;
            variableTree.tree.Configure(config);
            variableTree.tree.Train(trainingSet);

            ClassificationErrorDetails trainingErrorDetails = variableTree.tree.DatasetClassificationErrorDetails(trainingSet);
            ClassificationErrorDetails testErrorDetails = variableTree.tree.DatasetClassificationErrorDetails(testSet);
            writeErrorDetails("Training", trainingErrorDetails, stream);
            writeErrorDetails("Testing", testErrorDetails, stream);
            averageTrainingError += trainingErrorDetails.errorRate;
            averageTestError += testErrorDetails.errorRate;

            //cout << "Training error: " << tree.DescribeDatasetClassificationError(_trainingSet, true, testTreeInfo) << endl;
            //cout << "Test error: " << tree.DescribeDatasetClassificationError(_testSet, true, testTreeInfo) << endl;
        }

        //
        // Compute relative trees
        //
        for(auto relativeIterator = relatives.begin(); relativeIterator != relatives.end(); relativeIterator++)
        {
            for(UINT frameIndex = 0; frameIndex < frameCount - 1; frameIndex++)
            {
                _variableDataset.examples[frameIndex].output = Utility::BoolToInt(_frames[frameIndex + 1]->sourceFrame->observedInts[variableIndex] - _frames[frameIndex]->sourceFrame->observedInts[variableIndex] == *relativeIterator);
            }

            LearnerDataset trainingSet, testSet;
            _variableDataset.PartitionDataset(trainingSet, testSet, 0.7);

            stream << "Training tree for relative value: " << *relativeIterator << endl;

            GameVariableTree &variableTree = variableModel->trees[treeIndex++];
            variableTree.absolute = false;
            variableTree.value = *relativeIterator;
            variableTree.tree.Configure(config);
            variableTree.tree.Train(trainingSet);

            ClassificationErrorDetails trainingErrorDetails = variableTree.tree.DatasetClassificationErrorDetails(trainingSet);
            ClassificationErrorDetails testErrorDetails = variableTree.tree.DatasetClassificationErrorDetails(testSet);
            writeErrorDetails("Training", trainingErrorDetails, stream);
            writeErrorDetails("Testing", testErrorDetails, stream);

            averageTrainingError += trainingErrorDetails.errorRate;
            averageTestError += testErrorDetails.errorRate;

            //cout << "Training error: " << tree.DescribeDatasetClassificationError(_trainingSet, true, testTreeInfo) << endl;
            //cout << "Test error: " << tree.DescribeDatasetClassificationError(_testSet, true, testTreeInfo) << endl;
        }

        trainingFile << _game->ObservedIntNames()[variableIndex] << '\t' << averageTrainingError / (absolutes.size() + relatives.size()) << '\t' << averageTestError / (absolutes.size() + relatives.size()) << endl;
    }
}

const LearnerSprite& ModelLearner::SampleRandomSprite(int spriteID)
{
    int count = _spriteInstances[spriteID].size();
    if(count == 0)
    {
        return SampleRandomSprite(rand() % _parameters.totalSpriteCount);
    }
    else
    {
        return _spriteInstances[spriteID][rand() % count];
    }
}

void ModelLearner::TrainSpriteTrees(GameModel &model)
{
    const UINT frameCount = _frames.size();
    
    ofstream logFile("spriteTrainingLog.txt");
    ComposeStream stream(logFile, cout);

    DecisionTreeConfig config(50, 1, 100);

    model._spriteModels.resize(_parameters.totalSpriteCount);

    //
    // Learn decision trees for all variables
    //
    for(UINT spriteID = 0; spriteID < _parameters.totalSpriteCount; spriteID++)
    {
        stream << "Constructing sprite(X,Y) tree for sprite " << _game->Collection().GetSprite(spriteID).name << endl;

        //
        // Dataset is one-third positive examples, one-third negative examples, and one-third fully random examples (likely, including several empty ones)
        // The goal is to predict sprite existence in the next frame, so windows for examples are actually drawn from the *previous* frame(s).
        //
        for(int exampleIndex = 0; exampleIndex < int(_spriteDataset.examples.size()); exampleIndex++)
        {
            LearnerExample &curExample = _spriteDataset.examples[exampleIndex];
            Vec2i coord;
            int curFrameIndex;

            int exampleType = rand() % 3;
            if(exampleType == 0)
            {
                //
                // Positive example
                //
                const LearnerSprite &curSprite = SampleRandomSprite(spriteID);
                curFrameIndex = curSprite.frameIndex;
                coord = _frames[curFrameIndex]->sourceFrame->sprites[curSprite.spriteIndex].coord;
            }
            else if(exampleType == 1)
            {
                //
                // Negative example
                //
                const LearnerSprite &curSprite = SampleRandomSprite(rand() % _parameters.totalSpriteCount);
                curFrameIndex = curSprite.frameIndex;
                coord = _frames[curFrameIndex]->sourceFrame->sprites[curSprite.spriteIndex].coord;
            }
            else if(exampleType == 2)
            {
                //
                // Random example
                //
                const LearnerSprite &curSprite = SampleRandomSprite(rand() % _parameters.totalSpriteCount);
                coord = _frames[curSprite.frameIndex]->sourceFrame->sprites[curSprite.spriteIndex].coord;

                //
                // Pick a sample at this location but in a random frame
                //
                curFrameIndex = rand() % (_frames.size() - 5) + 5;
            }

            if(_frames[curFrameIndex]->sourceFrame->observedInts[0] == 1)
            {
                //
                // The 0th observed int is always the new game indicator. Learning based on the transition to a new frame is noisy, so we skip such examples.
                //
                exampleIndex--;
            }
            else
            {
                _frames[curFrameIndex - 1]->OverwriteDynamicFeatures(_parameters, coord.x / _parameters.gridScale, coord.y / _parameters.gridScale, curExample.input.data, true);
                curExample.output = _frames[curFrameIndex]->SpriteAtLocation(coord, spriteID);
            }
        }

        LearnerDataset trainingSet, testSet;
        _spriteDataset.PartitionDataset(trainingSet, testSet, 0.7);

        GameSpriteModel *newSpriteModel = new GameSpriteModel;
        
        newSpriteModel->spriteAtTree.Configure(config);
        newSpriteModel->spriteAtTree.Train(trainingSet);
        
        newSpriteModel->forceStasis = false;

        //if(spriteID == 0 || spriteID == 6 || spriteID == 7 || spriteID == 8) newSpriteModel->forceStasis = true;
        
        model._spriteModels[spriteID] = newSpriteModel;

        ClassificationErrorDetails trainingErrorDetails = newSpriteModel->spriteAtTree.DatasetClassificationErrorDetails(trainingSet);
        ClassificationErrorDetails testErrorDetails = newSpriteModel->spriteAtTree.DatasetClassificationErrorDetails(testSet);
        writeErrorDetails("Training", trainingErrorDetails, stream);
        writeErrorDetails("Testing", testErrorDetails, stream);
    }
}

void ModelLearner::ComputeVariableChanges(int variableIndex, set<int> &absolutes, set<int> &relatives)
{
    //
    // absolutes is the _parameters.maxAbsoluteVariableCount most frequently occuring values the variable takes
    // relatives is the set of differences after ignoring all those resolved by absolutes.
    //

    unordered_map<int, int> valueHistogram;
    const UINT frameCount = _frames.size();

    for(UINT frameIndex = 0; frameIndex < frameCount - 1; frameIndex++)
    {
        LearnerExample &curExample = _variableDataset.examples[frameIndex];
        GameModelFrame &curFrame = *_frames[frameIndex];
        GameModelFrame &nextFrame = *_frames[frameIndex + 1];

        int curValue = curFrame.sourceFrame->observedInts[variableIndex];
        int nextValue = nextFrame.sourceFrame->observedInts[variableIndex];

        if(valueHistogram.size() == 0 || curValue != nextValue)
        {
            if(valueHistogram.find(nextValue) == valueHistogram.end()) valueHistogram[nextValue] = 0;
            valueHistogram[nextValue]++;
        }
    }

    vector< pair<int, int> > sortedHistogram;

    for(auto it = valueHistogram.begin(); it != valueHistogram.end(); it++)
    {
        sortedHistogram.push_back(make_pair(it->first, it->second));
    }

    sort(sortedHistogram.begin(), sortedHistogram.end(), 
        [](const pair<int, int> & a, const pair<int, int> & b) -> bool
    {
        return a.second > b.second; 
    });

    absolutes.clear();
    for(int absoluteIndex = 0; absoluteIndex < _parameters.maxAbsoluteVariableCount && absoluteIndex < int(sortedHistogram.size()); absoluteIndex++)
    {
        absolutes.insert(sortedHistogram[absoluteIndex].first);
    }

    for(UINT frameIndex = 0; frameIndex < frameCount - 1; frameIndex++)
    {
        LearnerExample &curExample = _variableDataset.examples[frameIndex];
        GameModelFrame &curFrame = *_frames[frameIndex];
        GameModelFrame &nextFrame = *_frames[frameIndex + 1];

        int curValue = curFrame.sourceFrame->observedInts[variableIndex];
        int nextValue = nextFrame.sourceFrame->observedInts[variableIndex];
        int relativeValue = nextValue - curValue;
        
        if(absolutes.count(nextValue) == 0)
        {
            relatives.insert(relativeValue);
        }
    }
}
